import wx
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
import os
from Crypto.Random import get_random_bytes
import binascii
import webbrowser

def encrypt_file(input_path, output_path, key):
    try:
        with open(input_path, 'rb') as f_in:
            plaintext = f_in.read()
        
        cipher = AES.new(key, AES.MODE_CBC)
        ciphertext = cipher.encrypt(pad(plaintext, AES.block_size))
        
        # 确保输出路径有 .gcraesfile 后缀
        if not output_path.endswith('.gcraesfile'):
            output_path += '.gcraesfile'
        
        with open(output_path, 'wb') as f_out:
            f_out.write(cipher.iv)
            f_out.write(ciphertext)
        
        wx.MessageBox("文件加密成功！", "成功", wx.OK | wx.ICON_INFORMATION)
    except Exception as e:
        wx.MessageBox(f"加密过程中出现错误: {str(e)}", "错误", wx.OK | wx.ICON_ERROR)

def decrypt_file(input_path, output_path, key):
    try:
        with open(input_path, 'rb') as f_in:
            iv = f_in.read(16)
            ciphertext = f_in.read()
        
        cipher = AES.new(key, AES.MODE_CBC, iv)
        plaintext = unpad(cipher.decrypt(ciphertext), AES.block_size)
        
        # 移除 .gcraesfile 后缀
        if output_path.endswith('.gcraesfile'):
            output_path = output_path[:-len('.gcraesfile')]
        
        with open(output_path, 'wb') as f_out:
            f_out.write(plaintext)
        
        wx.MessageBox("文件解密成功！", "成功", wx.OK | wx.ICON_INFORMATION)
    except Exception as e:
        wx.MessageBox(f"解密过程中出现错误: {str(e)}", "错误", wx.OK | wx.ICON_ERROR)

class AesApp(wx.Frame):
    def __init__(self, parent=None, title="AES 文件加密解密 - LGCR837"):
        super().__init__(parent, title=title, size=(600, 300))
        
        panel = wx.Panel(self)
        
        # 添加作者信息
        author_label1 = wx.StaticText(panel, label="LGCR837制作")
        author_label2 = wx.StaticText(panel, label="Bilibili: https://space.bilibili.com/3493271839967884")
        
        self.input_file_path = wx.TextCtrl(panel, size=(300, -1))
        self.output_file_path = wx.TextCtrl(panel, size=(300, -1))
        self.key = wx.TextCtrl(panel, size=(300, -1))
        
        # 修改按钮文字
        input_button = wx.Button(panel, label="输入文件")
        output_button = wx.Button(panel, label="输出文件")
        auto_fill_button = wx.Button(panel, label="自动填充输出路径")
        generate_key_button = wx.Button(panel, label="随机")
        swap_button = wx.Button(panel, label="交换输入输出路径")
        encrypt_button = wx.Button(panel, label="加密")
        decrypt_button = wx.Button(panel, label="解密")
        
        input_button.Bind(wx.EVT_BUTTON, self.select_input_file)
        output_button.Bind(wx.EVT_BUTTON, self.select_output_file)
        auto_fill_button.Bind(wx.EVT_BUTTON, self.auto_fill_output)
        generate_key_button.Bind(wx.EVT_BUTTON, self.generate_random_key)
        swap_button.Bind(wx.EVT_BUTTON, self.swap_input_output)
        encrypt_button.Bind(wx.EVT_BUTTON, self.perform_encrypt)
        decrypt_button.Bind(wx.EVT_BUTTON, self.perform_decrypt)
        
        # 添加新按钮
        open_official_website_button = wx.Button(panel, label="打开官网")
        open_bilibili_button = wx.Button(panel, label="打开作者B站")
        
        open_official_website_button.Bind(wx.EVT_BUTTON, lambda event: webbrowser.open("https://cn.paydn.cn/app/fileaes.html"))
        open_bilibili_button.Bind(wx.EVT_BUTTON, lambda event: webbrowser.open("https://space.bilibili.com/3493271839967884"))
        
        # 调整 sizer 布局，增加一行
        sizer = wx.FlexGridSizer(rows=6, cols=3, hgap=10, vgap=10)
        sizer.Add(wx.StaticText(panel, label="输入文件:"), 0, wx.ALIGN_CENTER_VERTICAL)
        sizer.Add(self.input_file_path, 0, wx.EXPAND, 2)
        sizer.Add(input_button, 0, wx.EXPAND)
        
        sizer.Add(wx.StaticText(panel, label="输出文件:"), 0, wx.ALIGN_CENTER_VERTICAL)
        sizer.Add(self.output_file_path, 0, wx.EXPAND, 2)
        sizer.Add(output_button, 0, wx.EXPAND)
        
        sizer.Add(wx.StaticText(panel, label="密钥 (32 字节):"), 0, wx.ALIGN_CENTER_VERTICAL)
        sizer.Add(self.key, 0, wx.EXPAND, 2)
        sizer.Add(generate_key_button, 0, wx.EXPAND)
        
        sizer.Add(swap_button, 0, wx.EXPAND)
        # 将加密按钮跨两列
        sizer.Add(encrypt_button, 0, wx.EXPAND, 2)
        
        # 将自动填充按钮放在原来解密按钮的位置
        sizer.Add((0, 0))
        sizer.Add(auto_fill_button, 0, wx.EXPAND)
        sizer.Add(decrypt_button, 0, wx.EXPAND)
        
        # 添加新按钮到右下角
        sizer.Add((0, 0))
        sizer.Add(open_official_website_button, 0, wx.EXPAND)
        sizer.Add(open_bilibili_button, 0, wx.EXPAND)
        
        # 添加边框
        outer_sizer = wx.BoxSizer(wx.VERTICAL)
        outer_sizer.Add(author_label1, 0, wx.ALL, 5)
        outer_sizer.Add(author_label2, 0, wx.ALL, 5)
        outer_sizer.Add(sizer, 1, wx.ALL | wx.EXPAND, border=20)
        
        panel.SetSizer(outer_sizer)
        outer_sizer.Fit(self)
        
        self.Center()
        self.Show()
    
    def select_input_file(self, event):
        dialog = wx.FileDialog(self, "选择输入文件", style=wx.FD_OPEN)
        if dialog.ShowModal() == wx.ID_OK:
            self.input_file_path.SetValue(dialog.GetPath())
        dialog.Destroy()
    
    def select_output_file(self, event):
        dialog = wx.FileDialog(self, "选择输出文件", style=wx.FD_SAVE)
        if dialog.ShowModal() == wx.ID_OK:
            self.output_file_path.SetValue(dialog.GetPath())
        dialog.Destroy()
    
    def perform_encrypt(self, event):
        input_path = self.input_file_path.GetValue()
        output_path = self.output_file_path.GetValue()
        key_str = self.key.GetValue()
        
        if not input_path or not output_path or not key_str:
            wx.MessageBox("请填写所有字段！", "错误", wx.OK | wx.ICON_ERROR)
            return
        
        try:
            # 将十六进制字符串转换回字节类型
            key_bytes = binascii.unhexlify(key_str)
        except binascii.Error:
            wx.MessageBox("密钥必须是有效的字符串！", "错误", wx.OK | wx.ICON_ERROR)
            return
        
        encrypt_file(input_path, output_path, key_bytes)
    
    def perform_decrypt(self, event):
        input_path = self.input_file_path.GetValue()
        output_path = self.output_file_path.GetValue()
        key_str = self.key.GetValue()
        
        if not input_path or not output_path or not key_str:
            wx.MessageBox("请填写所有字段！", "错误", wx.OK | wx.ICON_ERROR)
            return
        
        try:
            # 将十六进制字符串转换回字节类型
            key_bytes = binascii.unhexlify(key_str)
        except binascii.Error:
            wx.MessageBox("密钥必须是有效的十六进制字符串！", "错误", wx.OK | wx.ICON_ERROR)
            return
        
        decrypt_file(input_path, output_path, key_bytes)
    
    def swap_input_output(self, event):
        temp = self.input_file_path.GetValue()
        self.input_file_path.SetValue(self.output_file_path.GetValue())
        self.output_file_path.SetValue(temp)

    def auto_fill_output(self, event):
        input_path = self.input_file_path.GetValue()
        if input_path:
            file_dir, file_name = os.path.split(input_path)
            file_base, _ = os.path.splitext(file_name)
            output_path = os.path.join(file_dir, file_base + '.gcraesfile')
            self.output_file_path.SetValue(output_path)

    def generate_random_key(self, event):
        # 生成 32 字节的随机密钥
        key = get_random_bytes(32)
        # 将密钥转换为十六进制字符串显示给用户
        key_hex = binascii.hexlify(key).decode('utf-8')
        self.key.SetValue(key_hex)

if __name__ == "__main__":
    app = wx.App()
    frame = AesApp()
    app.MainLoop()